#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>

#define MATRIX_SIZE 100

typedef struct {
	int row;
	int col;
	float mData[MATRIX_SIZE];
} MATRIX_DATA;

MATRIX_DATA data;

MPI_Datatype NEW_MPI_TYPE;

void create_new_type(MPI_Datatype *pnewtype){
	// count of length each block 
	int block_lens[3] = {1, 1, MATRIX_SIZE};	
	
	// Displacement for each item a, b, n
	MPI_Aint disps[3];
	MPI_Aint start_addr, next_addr;
	disps[0] = 0;
	MPI_Address( &data.row, &start_addr );
	MPI_Address( &data.col, &next_addr );
	disps[1] = next_addr - start_addr;
	MPI_Address( &data.mData, &next_addr );
	disps[2] = next_addr - start_addr;
	
	// Define list of MPI_Datatype for each item in struct
	MPI_Datatype typelist[3] = {MPI_INT, MPI_INT, MPI_FLOAT};
	
	// Build the derived MPI struct
	MPI_Type_struct( 	3, 		/* Có 3 khoi: a, b va n */
				 block_lens, 	/* Số phan tu tren tung khoi */
				 disps,		/* Độ dời từng khối */
				 typelist, 	/* Danh sach cac kieu du lieu */
				 pnewtype	/* Kieu MPI moi */				  
				 );
					 
	MPI_Type_commit( pnewtype );
}

int main(int argc, char* argv[]) {
	int rank, size, i, j, k, len; 
	MPI_Init( &argc, &argv );
	MPI_Comm_size(MPI_COMM_WORLD, &size);
	MPI_Comm_rank( MPI_COMM_WORLD, &rank );
	MATRIX_DATA data1, data2;
	int offset[size], r_offset; 
	FILE *f1, *f2;
	if (rank == 0) {
		
		// get data 
		f1 = fopen("input1.txt","r");
		fscanf(f1, "%d %d", &data1.row, &data1.col);
		
		f2 = fopen("input2.txt","r");
		fscanf(f2, "%d %d", &data2.row, &data2.col);
	
		//data1.mData = (float*)malloc(data1.row*data1.col*sizeof(float));
		for(i = 0; i < data1.row * data1.col; i++)
			fscanf(f1, "%f", &data1.mData[i]);
	
		//data2.mData = (float*)malloc(data2.row*data2.col*sizeof(float));
		for(i = 0; i < data2.row * data2.col; i++)
			fscanf(f2, "%f", &data2.mData[i]);
	
	}
	
	create_new_type( &NEW_MPI_TYPE);
	
	MPI_Barrier( MPI_COMM_WORLD );
	
	MPI_Bcast( &data1, 1, NEW_MPI_TYPE, 0, MPI_COMM_WORLD );
	MPI_Bcast(&data2, 1, NEW_MPI_TYPE, 0, MPI_COMM_WORLD);
	printf("size = %d\n", size);
	len = data1.row / size;
	if(rank == 0){
		// cal offset
		offset[0] = 0;
		for (i = 1; i < size; i++){
			offset[i] = offset[i-1]+len;
			//printf("offset = %d\n", offset[i]);
		}
	}
	MPI_Scatter(offset,1, MPI_INT, &r_offset, 1, MPI_INT, 0, MPI_COMM_WORLD); 
	printf("offset = %d\n", r_offset);
	float result[len* data2.col], m_result[data1.row*data2.col];
	
	for(i = 0; i < len; i++){
		for(j = 0;j < data2.col; j++){
			// row = i + offset
			result[i*data2.col+j] = 0;
			for(k = 0; k < data1.col; k++)
				result[i* data2.col + j] += data1.mData[(i+r_offset)*data1.col+ k]* data2.mData[k*data2.col+j];
		}
	}
	MPI_Gather(result, len*data2.col, MPI_FLOAT, m_result, len*data2.col, MPI_FLOAT, 0, MPI_COMM_WORLD);
	
	// print result
	if(rank == 0){
		for(i = 0; i < data1.row; i++){
			for(j = 0; j < data2.col; j++)
				printf("%.1f\t", m_result[i*data2.col+j]);
			printf("\n");
				
		}
	}		
	MPI_Finalize();
	return 0;
} // End main
